'''
Author: didiplus
Description: 
Date: 2025-08-21 11:23:37
LastEditors: didiplus
LastEditTime: 2025-08-21 15:52:40
FilePath: /WeChatBot/app/core/database.py
Version: 1.0
'''
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker, Session
from sqlalchemy.ext.declarative import declarative_base
from typing import Generator
from contextlib import contextmanager
from .settings import settings
from models import User, KeywordRuleModel, Base


class Database:
    def __init__(self):
        self.engine = None
        self.SessionLocal = None
    
    def init_engine(self):
        """初始化数据库引擎"""
        self.engine = create_engine(
            settings.DATABASE_URL,
            pool_size=10,
            max_overflow=20,
            pool_timeout=30,
            pool_recycle=1800,
            pool_pre_ping=True,
            echo=settings.DEBUG,
            connect_args={
                "connect_timeout": 10,
            }
        )
        
        self.SessionLocal = sessionmaker(
            autocommit=False,
            autoflush=False,
            bind=self.engine
        )
    
    def get_session(self) -> Session:
        """获取单个会话（需要手动关闭）"""
        if self.SessionLocal is None:
            self.init_engine()
        return self.SessionLocal()
        
    @contextmanager
    def get_db_session(self) -> Generator[Session, None, None]:
        """获取数据库会话上下文管理器"""
        if self.SessionLocal is None:
            self.init_engine()
        
        db = self.SessionLocal()
        try:
            yield db
        finally:
            db.close()

# 全局实例
database = Database()

# 修改这里的实现
def get_db():
    """
    FastAPI依赖注入使用的数据库会话生成器
    """
    if database.SessionLocal is None:
        database.init_engine()
    
    db = database.SessionLocal()
    try:
        yield db
    finally:
        db.close()

def get_db_session() -> Session:
    """直接获取数据库会话（需要手动管理）"""
    return database.get_session()

def init_db():
    """
    初始化数据库，创建所有表
    """
    if database.engine is None:
        database.init_engine()
    
    try:
        Base.metadata.create_all(bind=database.engine)
        print("数据库表创建成功")
        return True
    except Exception as e:
        print(f"创建数据库表时出错: {e}")
        raise